{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OojBjVLB-XdL"
   },
   "source": [
    "# 导入必要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.36.2\n",
      "0.19.4\n",
      "1.26.4\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.12/site-packages/transformers/utils/generic.py:441: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.\n",
      "  _torch_pytree._register_pytree_node(\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "import huggingface_hub\n",
    "import numpy\n",
    "\n",
    "print(transformers.__version__)\n",
    "print(huggingface_hub.__version__)\n",
    "print(numpy.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "executionInfo": {
     "elapsed": 15616,
     "status": "ok",
     "timestamp": 1758722658663,
     "user": {
      "displayName": "sy l",
      "userId": "08128253523941293274"
     },
     "user_tz": -480
    },
    "id": "t0nJ9bNo8R61"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.12/site-packages/transformers/utils/generic.py:309: FutureWarning: `torch.utils._pytree._register_pytree_node` is deprecated. Please use `torch.utils._pytree.register_pytree_node` instead.\n",
      "  _torch_pytree._register_pytree_node(\n"
     ]
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import torch\n",
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "from transformers import BertForSequenceClassification, BertTokenizer, Trainer, TrainingArguments\n",
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "executionInfo": {
     "elapsed": 7,
     "status": "ok",
     "timestamp": 1758722658726,
     "user": {
      "displayName": "sy l",
      "userId": "08128253523941293274"
     },
     "user_tz": -480
    },
    "id": "vITxToWs8klk"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 设备配置\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DIGDsVWF8ygr"
   },
   "source": [
    "# 加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading Model from https://www.modelscope.cn to directory: /root/autodl-tmp/models/pretrained/google-bert/bert-base-chinese\n"
     ]
    }
   ],
   "source": [
    "from modelscope import snapshot_download\n",
    "\n",
    "model_dir = snapshot_download('google-bert/bert-base-chinese', cache_dir='/root/autodl-tmp/models/pretrained')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "editable": true,
    "id": "kNXA4PNB803z",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at /root/autodl-tmp/models/pretrained/google-bert/bert-base-chinese and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from modelscope import AutoModelForMaskedLM, AutoTokenizer\n",
    "from modelscope import snapshot_download\n",
    "\n",
    "\n",
    "model_dir = '/root/autodl-tmp/models/pretrained/google-bert/bert-base-chinese'\n",
    "# 对应MLM任务\n",
    "# model = AutoModelForMaskedLM.from_pretrained(model_dir)\n",
    "model = BertForSequenceClassification.from_pretrained(\n",
    "    model_dir,\n",
    "    num_labels=2  # 二分类\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BertForSequenceClassification(\n",
      "  (bert): BertModel(\n",
      "    (embeddings): BertEmbeddings(\n",
      "      (word_embeddings): Embedding(21128, 768, padding_idx=0)\n",
      "      (position_embeddings): Embedding(512, 768)\n",
      "      (token_type_embeddings): Embedding(2, 768)\n",
      "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "      (dropout): Dropout(p=0.1, inplace=False)\n",
      "    )\n",
      "    (encoder): BertEncoder(\n",
      "      (layer): ModuleList(\n",
      "        (0-11): 12 x BertLayer(\n",
      "          (attention): BertAttention(\n",
      "            (self): BertSelfAttention(\n",
      "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "            (output): BertSelfOutput(\n",
      "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
      "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "              (dropout): Dropout(p=0.1, inplace=False)\n",
      "            )\n",
      "          )\n",
      "          (intermediate): BertIntermediate(\n",
      "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
      "            (intermediate_act_fn): GELUActivation()\n",
      "          )\n",
      "          (output): BertOutput(\n",
      "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
      "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
      "            (dropout): Dropout(p=0.1, inplace=False)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (pooler): BertPooler(\n",
      "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
      "      (activation): Tanh()\n",
      "    )\n",
      "  )\n",
      "  (dropout): Dropout(p=0.1, inplace=False)\n",
      "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O1XPLoax9fiV"
   },
   "source": [
    "# LoRA配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "8h_Xa8tT9ic5"
   },
   "outputs": [],
   "source": [
    "# lora_config = LoraConfig(\n",
    "#     task_type=TaskType.SEQ_CLS,\n",
    "#     r=8,\n",
    "#     lora_alpha=32,\n",
    "#     target_modules=[\"query\", \"key\", \"value\"],\n",
    "#     lora_dropout=0.1,\n",
    "#     bias=\"none\"\n",
    "# )\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.SEQ_CLS,\n",
    "    r=8,\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"attention.self.query\", \"attention.self.key\", \"attention.self.value\"],\n",
    "    lora_dropout=0.1,\n",
    "    bias=\"none\"\n",
    ")\n",
    "\n",
    "# 应用LoRA\n",
    "model = get_peft_model(model, lora_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "XXZJw75Y9rZ7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "可训练参数总数: 443906\n"
     ]
    }
   ],
   "source": [
    "# 输出可训练参数数量\n",
    "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(f\"可训练参数总数: {trainable_params}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dGRYcJph9uy1"
   },
   "source": [
    "# 加载数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "KeyboardInterrupt\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset, DownloadConfig  # 导入 DownloadConfig 类\n",
    "\n",
    "# 1. 创建 DownloadConfig 实例，并配置下载参数\n",
    "download_config = DownloadConfig()\n",
    "download_config.resume_download = True  # 开启断点续传\n",
    "download_config.max_retries = 5         # 下载失败自动重试 5 次\n",
    "download_config.speed_limit = 3 * 1024 * 1024  # 限速 3MB/s（单位：B/s）\n",
    "\n",
    "# 2. 加载数据集（此时 download_config 是正确的类实例）\n",
    "dataset = load_dataset(\n",
    "    path=\"lansinuote/ChnSentiCorp\",  # 数据集名称\n",
    "    cache_dir=\"/root/autodl-tmp/data/raw/ChnSentiCorp\",  # 自定义缓存目录\n",
    "    token=\"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\",  # 你的 HF Token\n",
    "    download_config=download_config  # 传入 DownloadConfig 实例（而非字典）\n",
    ")\n",
    "\n",
    "# 验证加载成功\n",
    "print(\"数据集下载完成！\")\n",
    "print(f\"数据集拆分：{dataset.keys()}\")  # 输出 train/validation/test 等拆分\n",
    "print(f\"训练集样本数：{len(dataset['train'])}\")\n",
    "print(f\"第一个样本：\\n{dataset['train'][0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "id": "4GN1st5y9xS3",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# 加载本地数据集并裁剪样本\n",
    "dataset = load_dataset(\"/root/autodl-tmp/data/raw/ChnSentiCorp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nrZV0dKf99fr"
   },
   "outputs": [],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WkvPdBvl-ELW"
   },
   "outputs": [],
   "source": [
    "# 分词函数\n",
    "def tokenize(batch):\n",
    "    return tokenizer(batch['text'], padding='max_length', truncation=True, max_length=128)\n",
    "\n",
    "\n",
    "# 数据集分词\n",
    "tokenized_datasets = dataset.map(tokenize, batched=True)\n",
    "tokenized_datasets.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CaChzGD_-IAC"
   },
   "outputs": [],
   "source": [
    "# 计算评估指标\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels = eval_pred\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='binary')\n",
    "    acc = accuracy_score(labels, predictions)\n",
    "    return {\"accuracy\": acc, \"f1\": f1, \"precision\": precision, \"recall\": recall}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R1lsz_Wg-RB-"
   },
   "source": [
    "# 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rHrSzMRV-LQH"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# 训练参数\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='/tmp/code/bert_fine_tuning/models/fine_tuned/lora_results',\n",
    "    num_train_epochs=3,\n",
    "    per_device_train_batch_size=16,\n",
    "    per_device_eval_batch_size=16,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=100,\n",
    "    save_steps=100,\n",
    "    logging_steps=50,\n",
    "    learning_rate=2e-4,\n",
    "    weight_decay=0.01,\n",
    "    save_total_limit=1,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"accuracy\"\n",
    ")\n",
    "\n",
    "# 初始化Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_datasets['train'],\n",
    "    eval_dataset=tokenized_datasets['test'],\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics\n",
    ")\n",
    "\n",
    "# 开始训练\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型保存与加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import PeftModel\n",
    "\n",
    "# 保存LoRA模型\n",
    "save_dir = \"/tmp/code/bert_fine_tuning/models/fine_tuned/version1\"\n",
    "model.save_pretrained(save_dir)\n",
    "tokenizer.save_pretrained(save_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import BertForSequenceClassification, BertTokenizer\n",
    "from peft import PeftModel\n",
    "\n",
    "# 模型路径\n",
    "model_dir = \"/tmp/code/bert_fine_tuning/models/pretrained/google-bert/bert-base-chinese\"\n",
    "save_dir = \"/tmp/code/bert_fine_tuning/models/fine_tuned/version1\"\n",
    "\n",
    "# 加载基础模型\n",
    "base_model = BertForSequenceClassification.from_pretrained(\n",
    "    model_dir,\n",
    "    num_labels=2\n",
    ")\n",
    "\n",
    "# 加载LoRA权重\n",
    "model = PeftModel.from_pretrained(base_model, save_dir)\n",
    "\n",
    "# 指定设备\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device)\n",
    "model.eval()\n",
    "\n",
    "# 加载分词器\n",
    "tokenizer = BertTokenizer.from_pretrained(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 单条推理函数\n",
    "def test(text):\n",
    "    if not text.strip():\n",
    "        print(\"输入为空，无法分类\")\n",
    "        return\n",
    "    # 编码\n",
    "    inputs = tokenizer(\n",
    "        text,\n",
    "        padding=True,\n",
    "        truncation=True,\n",
    "        max_length=128,\n",
    "        return_tensors=\"pt\"\n",
    "    )\n",
    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
    "    \n",
    "    # 推理\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**inputs)\n",
    "        logits = outputs.logits\n",
    "        probs = torch.softmax(logits, dim=-1)\n",
    "    \n",
    "    predicted_class = torch.argmax(probs, dim=-1).item()\n",
    "    print(f\"文本: {text}\")\n",
    "    print(\"分类结果:\", predicted_class)\n",
    "    print(\"类别概率:\", probs.tolist())\n",
    "    print(\"-\" * 40)\n",
    "\n",
    "\n",
    "# 批量测试\n",
    "texts = [\"这个电影真的很好看！\", \"\", \"这部电影太无聊了\", \"这部电影一般般\", \"还可以\"]\n",
    "for t in texts:\n",
    "    test(t)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyN6Vq9WdQG4QaI5+TYDb6Fm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
